import cv2
import threading
import queue
import numpy as np


class RTSPCaptureThread(threading.Thread):
    def __init__(self, rtsp_url, frame_queue):
        super(RTSPCaptureThread, self).__init__()
        self.rtsp_url = rtsp_url
        self.cap = cv2.VideoCapture(self.rtsp_url)
        self.frame_queue = frame_queue
        self.running = True

    def run(self):
        while self.running:
            ret, frame = self.cap.read()
            if not ret:
                print(f"Failed to grab frame from {self.rtsp_url}")
                continue
            self.frame_queue.put(frame)
        self.cap.release()

    def stop(self):
        self.running = False


class DisplayThread(threading.Thread):
    def __init__(self, frame_queues):
        super(DisplayThread, self).__init__()
        self.frame_queues = frame_queues
        self.running = True

    def run(self):
        while self.running:
            n = len(self.frame_queues)
            frams = []
            for idx in range(n):
                tmp_fram = self.frame_queues[idx].get()
                tmp_fram = cv2.resize(tmp_fram, (640, 360))
                frams.append(tmp_fram)
            # 合并摄像头帧
            combined_frame_h1 = np.hstack(frams[0:n//2])
            combined_frame_h2 = np.hstack(frams[n//2:n])
            combined_frame = np.vstack((combined_frame_h1, combined_frame_h2))

            cv2.imshow("RTSP Stream", combined_frame)
            if cv2.waitKey(5) & 0xFF == ord('q'):
                self.stop()

    def stop(self):
        self.running = False
        cv2.destroyAllWindows()


if __name__ == "__main__":

    rtsp_urls = [
        'rtsp://admin:jk123456@192.168.20.27:554/h264/ch1/main/av_stream',
        'rtsp://admin:jk123456@192.168.20.28:554/h264/ch1/main/av_stream',
        'rtsp://admin:jk123456@192.168.20.29:554/h264/ch1/main/av_stream',
        'rtsp://admin:jk123456@192.168.20.59:554/h264/ch1/main/av_stream',
        'rtsp://admin:jk123456@192.168.20.61:554/h264/ch1/main/av_stream',
        'rtsp://admin:jk123456@192.168.20.63:554/h264/ch1/main/av_stream',
        # 'rtsp://admin:jk123456@192.168.20.71:554/h264/ch1/main/av_stream',
        # 'rtsp://admin:jk123456@192.168.20.72:554/h264/ch1/main/av_stream',
    ]

    # 创建一个帧队列列表，每个流一个队列
    frame_queues = [queue.Queue(maxsize=5) for _ in range(len(rtsp_urls))]

    # 创建空列表，存放每个流的捕获和显示线程
    capture_threads = []
    display_threads = []

    # 启动每个RTSP取流线程
    for idx, rtsp_url in enumerate(rtsp_urls):
        thread = RTSPCaptureThread(rtsp_url, frame_queues[idx])
        thread.start()
        capture_threads.append(thread)

    display_threads = DisplayThread(frame_queues)
    display_threads.start()

    # 这里，线程会持续运行，直到按下'q'键关闭一个显示窗口。你也可以设置其他退出机制。
    for capture_thread in capture_threads:
        print(f"{capture_thread}线程开始运行")
        capture_thread.join()

    display_threads.join()

# 显示帧数
    # print(threads)
    # # 按需运行一段时间（例如10秒），然后停止所有线程
    # import time
    # time.sleep(10)  # 运行10秒

    # for thread in threads:
    #     thread.stop()
    #     thread.join()

    # # 输出每个线程捕获的帧数
    # for thread in threads:
    #     print(
    #         f"Stream {thread.rtsp_url} captured {len(thread.frames)} frames.")
